package com.exm.collectcodenew.web.controller;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.util.HashMap;
import java.util.Map;

public class AgentContentAggregator {

    private static final ObjectMapper objectMapper = new ObjectMapper();
    private final Map<String, StringBuilder> agentContentMap = new HashMap<>();

    /**
     * 处理单条 SSE 数据行
     */
    public void processDataLine(String dataLine) {
        if (dataLine == null || dataLine.trim().isEmpty()) {
            return;
        }

        // 移除 "data: " 前缀并trim
        String jsonData = dataLine.startsWith("data: ") ? dataLine.substring(6).trim() : dataLine.trim();

        if ("[DONE]".equals(jsonData)) {
            System.out.println("流式传输结束");
            return;
        }

        try {
            JsonNode jsonNode = objectMapper.readTree(jsonData);
            processJsonNode(jsonNode);
        } catch (Exception e) {
            System.err.println("解析 JSON 失败: " + jsonData);
            e.printStackTrace();
        }
    }

    /**
     * 处理 JSON 节点数据
     */
    private void processJsonNode(JsonNode jsonNode) {
        if (!jsonNode.has("agent") || !jsonNode.has("content")) {
            return;
        }

        String agent = jsonNode.get("agent").asText();
        String content = jsonNode.get("content").asText();

        // 将内容按 agent 类型拼接
        appendContent(agent, content);
    }

    /**
     * 将内容按 agent 类型拼接
     */
    private void appendContent(String agent, String content) {
        agentContentMap.compute(agent, (key, existingBuilder) -> {
            if (existingBuilder == null) {
                existingBuilder = new StringBuilder();
            }
            existingBuilder.append(content);
            return existingBuilder;
        });
    }

    /**
     * 获取最终的拼接结果
     */
    public Map<String, String> getAggregatedResult() {
        Map<String, String> result = new HashMap<>();
        agentContentMap.forEach((agent, stringBuilder) -> {
            result.put(agent, stringBuilder.toString());
        });
        return result;
    }

    /**
     * 获取原始的 StringBuilder 映射（用于进一步操作）
     */
    public Map<String, StringBuilder> getAgentContentMap() {
        return agentContentMap;
    }

    /**
     * 清空当前内容
     */
    public void clear() {
        agentContentMap.clear();
    }

    /**
     * 处理示例数据
     */
    public void processExampleData() {
        // 您提供的示例数据
        String[] exampleData = {
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"coordinator\",\"created\":1642428800,\"content\":\"你好\",\"index\":0,\"finish_reason\":null}",
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"coordinator\",\"created\":1642428800,\"content\":\"，\",\"index\":0,\"finish_reason\":null}",
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"coordinator\",\"created\":1642428800,\"content\":\"有\",\"index\":0,\"finish_reason\":null}",
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"coordinator\",\"created\":1642428800,\"content\":\"什么\",\"index\":0,\"finish_reason\":null}",
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"coordinator\",\"created\":1642428800,\"content\":\"可以\",\"index\":0,\"finish_reason\":null}",
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"planner\",\"created\":1642428800,\"content\":\"帮\",\"index\":0,\"finish_reason\":null}",
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"planner\",\"created\":1642428800,\"content\":\"您\",\"index\":0,\"finish_reason\":null}",
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"planner\",\"created\":1642428800,\"content\":\"的\",\"index\":0,\"finish_reason\":null}",
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"planner\",\"created\":1642428800,\"content\":\"吗\",\"index\":0,\"finish_reason\":null}",
                "data: {\"id\":\"chatcmpl-9ABC123\",\"agent\":\"planner\",\"created\":1642428800,\"content\":\"？\",\"index\":0,\"finish_reason\":null}"
        };

        System.out.println("开始处理流数据...");

        for (String dataLine : exampleData) {
            processDataLine(dataLine);
        }

        // 获取并打印结果
        Map<String, String> result = getAggregatedResult();
        System.out.println("\n=== 聚合结果 ===");
        result.forEach((agent, content) -> {
            System.out.println("Agent: " + agent);
            System.out.println("拼接内容: " + content);
            System.out.println("内容长度: " + content.length());
            System.out.println("---");
        });
    }

    public static void main(String[] args) {
        AgentContentAggregator aggregator = new AgentContentAggregator();
        aggregator.processExampleData();
    }
}